import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, LSTM, Dense, Attention, TimeDistributed
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Dropout

# 设置中文显示
plt.rcParams["font.family"] = ["SimSun"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 加载数据
try:
    stock_data = pd.read_csv('H:\\csv_shares_day\\000001-19900101-20211231.csv',encoding="gbk")
    stock_data.columns = ["Date", "code", "name", "Close", "High", "Low", "Open", "pre_close", "subval",
                          "subrate",
                          "turnrate",
                          "Volume", "Money", "all_earn", "exchage_earn"]
    stock_data.set_index(['Date'], inplace=True)  # 修正索引列名
    dfall = stock_data.sort_values(by=["Date"])
    data = dfall.loc["2016-01-01":"2021-12-31"]
    data.set_index('Date', inplace=True)
except Exception as e:
    print(f"数据加载错误: {e}")
    # 生成示例数据用于测试
    dates = pd.date_range(start='2023-01-01', periods=365)
    values = np.random.randn(365).cumsum() + 100
    data = pd.DataFrame({'Close': values}, index=dates)

# 仅使用'Close'价格进行预测
data = data[['Close']]

# 数据标准化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(data)

# 创建数据集函数
def create_dataset(data, look_back=1):
    X, Y = [], []
    for i in range(len(data) - look_back):
        X.append(data[i:(i + look_back), 0])
        Y.append(data[i + look_back, 0])
    return np.array(X), np.array(Y)

look_back = 30  # 修改时间窗口大小为30
X, Y = create_dataset(scaled_data, look_back)
X = np.reshape(X, (X.shape[0], X.shape[1], 1))  # 重塑为 [样本数, 时间步, 特征数]

# 构建模型
def build_model(input_shape):
    inputs = Input(shape=input_shape)
    lstm_out = LSTM(50, return_sequences=True)(inputs)
    attention_out = Attention()([lstm_out, lstm_out])
    dropout_out = TimeDistributed(Dropout(0.2))(attention_out)
    outputs = TimeDistributed(Dense(1))(dropout_out)

    model = Model(inputs=inputs, outputs=outputs)
    model.compile(optimizer='adam', loss='mean_squared_error')
    return model

model = build_model((look_back, 1))

# 分割数据集
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)
Y_train = Y_train.reshape(-1, 1, 1)

# 训练模型
history = model.fit(X_train, Y_train, epochs=1000, batch_size=32, validation_split=0.2, verbose=1)

# 可视化训练历史
plt.figure(figsize=(10, 6))
plt.plot(history.history['loss'], label='训练损失')
plt.plot(history.history['val_loss'], label='验证损失')
plt.title('模型训练损失')
plt.xlabel('轮次')
plt.ylabel('损失')
plt.legend()
plt.grid(True)
plt.show()

# 对测试数据进行预测
test_predict = model.predict(X_test)
# 只取最后一个时间步的预测值
test_predict = test_predict[:, -1, 0].reshape(-1, 1)
# 反归一化
test_predict = scaler.inverse_transform(test_predict)
# 反归一化真实值
Y_test_actual = scaler.inverse_transform(Y_test.reshape(-1, 1))

# 可视化测试数据预测结果
plt.figure(figsize=(12, 6))
plt.plot(Y_test_actual, label='实际股价')
plt.plot(test_predict, label='预测股价')
plt.title('测试数据股价预测与实际对比')
plt.xlabel('样本')
plt.ylabel('股价')
plt.legend()
plt.grid(True)
plt.show()

# 预测未来5天的股价（修改为预测5天）
last_sequence = scaled_data[-look_back:].reshape(1, look_back, 1)
future_predictions = []

for _ in range(5):  # 修改为预测5天
    next_day = model.predict(last_sequence, verbose=0)
    predicted_value = next_day[0, -1, 0]
    future_predictions.append(predicted_value)
    last_sequence = np.roll(last_sequence, -1, axis=1)
    last_sequence[0, -1, 0] = predicted_value

# 反归一化并打印预测结果
future_predictions = np.array(future_predictions).reshape(-1, 1)
future_predictions = scaler.inverse_transform(future_predictions)

print("未来5天的预测股价:")
for i, price in enumerate(future_predictions, 1):
    print(f"第{i}天: {price[0]:.2f}")

# 可视化预测结果
plt.figure(figsize=(14, 7))
plt.plot(data.index, data['Close'], label='历史股价', color='blue')

last_date = data.index[-1]
future_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=5)  # 生成5天日期

plt.plot(future_dates, future_predictions, 'r--o', label='预测股价')
plt.plot([data.index[-1], future_dates[0]], [data['Close'].iloc[-1], future_predictions[0, 0]], 'k--', alpha=0.5)

plt.title('股价预测')
plt.xlabel('日期')
plt.ylabel('价格')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

# 计算测试集预测的评估指标
mse = mean_squared_error(Y_test_actual, test_predict)
rmse = np.sqrt(mse)
mae = mean_absolute_error(Y_test_actual, test_predict)
mape = np.mean(np.abs((Y_test_actual - test_predict) / Y_test_actual)) * 100
r2 = r2_score(Y_test_actual, test_predict)

# 打印评估指标
print("\n模型评估指标:")
print(f"均方误差(MSE): {mse:.4f}")
print(f"均方根误差(RMSE): {rmse:.4f}")
print(f"平均绝对误差(MAE): {mae:.4f}")
print(f"平均绝对百分比误差(MAPE): {mape:.2f}%")
print(f"决定系数(R²): {r2:.4f}")

# 计算方向准确率（预测涨跌方向的准确率）
direction_accuracy = np.mean((Y_test_actual[1:] - Y_test_actual[:-1]) *
                            (test_predict[1:] - test_predict[:-1]) > 0) * 100
print(f"方向预测准确率: {direction_accuracy:.2f}%")

# 保存模型
model.save('stock_prediction_model.keras')
print("模型已保存到 'stock_prediction_model.keras'")